import abc
from collections import OrderedDict

import gtimer as gt

from rlkit.core import logger, eval_util
from rlkit.data_management.replay_buffer import ReplayBuffer
from rlkit.samplers.data_collector import DataCollector

from torch.utils.tensorboard import SummaryWriter


def _get_epoch_timings():
    times_itrs = gt.get_times().stamps.itrs
    times = OrderedDict()
    epoch_time = 0
    for key in sorted(times_itrs):
        time = times_itrs[key][-1]
        epoch_time += time
        times['time/{} (s)'.format(key)] = time
    times['time/epoch (s)'] = epoch_time
    times['time/total (s)'] = gt.get_times().total
    return times


class BaseRLAlgorithm(object, metaclass=abc.ABCMeta):
    def __init__(
            self,
            trainer,
            exploration_env,
            evaluation_env,
            exploration_data_collector: DataCollector,
            evaluation_data_collector: DataCollector,
            replay_buffer: ReplayBuffer,
            expl_replay_buffer: ReplayBuffer = None,  # for jsrl
    ):
        self.trainer = trainer
        self.expl_env = exploration_env
        self.eval_env = evaluation_env
        self.expl_data_collector = exploration_data_collector
        self.eval_data_collector = evaluation_data_collector
        self.replay_buffer = replay_buffer
        self.expl_replay_buffer = expl_replay_buffer
        self._start_epoch = 0

        self.post_epoch_funcs = []

        self.writer = SummaryWriter(log_dir=logger.get_snapshot_dir())  # 初始化 TensorBoard writer

    def train(self, start_epoch=0):
        self._start_epoch = start_epoch
        self._train()

    def _train(self):
        """
        Train model.
        """
        raise NotImplementedError('_train must implemented by inherited class')

    def _begin_epoch(self, epoch):
        pass

    # _end_epoch
    def _end_epoch(self, epoch):
        snapshot = self._get_snapshot()
        logger.save_itr_params(epoch, snapshot)
        gt.stamp('saving')
        self._log_stats(epoch)

        self.expl_data_collector.end_epoch(epoch)
        self.eval_data_collector.end_epoch(epoch)
        self.replay_buffer.end_epoch(epoch)
        if self.expl_replay_buffer is not None:
            self.expl_replay_buffer.end_epoch(epoch)
        self.trainer.end_epoch(epoch)

        for post_epoch_func in self.post_epoch_funcs:
            post_epoch_func(self, epoch)

    def _get_snapshot(self):
        snapshot = {}
        for k, v in self.trainer.get_snapshot().items():
            snapshot['trainer/' + k] = v
        for k, v in self.expl_data_collector.get_snapshot().items():
            snapshot['exploration/' + k] = v
        for k, v in self.eval_data_collector.get_snapshot().items():
            snapshot['evaluation/' + k] = v
        for k, v in self.replay_buffer.get_snapshot().items():
            snapshot['replay_buffer/' + k] = v
        return snapshot

    def _log_stats(self, epoch):
        logger.log("Epoch {} finished".format(epoch), with_timestamp=True)
        logger.record_dict({"epoch": epoch})

        """
        Replay Buffer
        """
        replay_stats = self.replay_buffer.get_diagnostics()
        logger.record_dict(replay_stats, prefix='replay_buffer/')
        # 记录 replay buffer size 到 TensorBoard
        self.writer.add_scalar('replay_buffer/offline_size', replay_stats.get('size', 0), epoch)

        if self.expl_replay_buffer is not None:
            expl_replay_stats = self.expl_replay_buffer.get_diagnostics()
            logger.record_dict(expl_replay_stats, prefix='expl_replay_buffer/')
            # 记录 exploration replay buffer size 到 TensorBoard
            self.writer.add_scalar('replay_buffer/online_size', expl_replay_stats.get('size', 0), epoch)

        """
        Trainer
        """
        trainer_stats = self.trainer.get_diagnostics()
        # logger.record_dict(trainer_stats, prefix='trainer/')
        
        # 记录训练损失到 TensorBoard
        self.writer.add_scalar('trainer/QF1 Loss', trainer_stats.get('QF1 Loss', 0), epoch)
        self.writer.add_scalar('trainer/QF2 Loss', trainer_stats.get('QF2 Loss', 0), epoch)
        self.writer.add_scalar('trainer/Policy Loss', trainer_stats.get('Policy Loss', 0), epoch)
        self.writer.add_scalar('trainer/VF Loss', trainer_stats.get('VF Loss', 0), epoch)

        """
        Exploration
        """
        # path length step 等
        expl_path_info = self.expl_data_collector.get_diagnostics()
        # logger.record_dict(expl_path_info, prefix='expl/')
        self.writer.add_scalar('expl/path length Mean', expl_path_info.get('path length Mean', 0), epoch)
        # xxx
        expl_paths = self.expl_data_collector.get_epoch_paths()
        # if hasattr(self.expl_env, 'get_diagnostics'):
        #     logger.record_dict(self.expl_env.get_diagnostics(expl_paths), prefix='expl/')
        # Returns Reward 等
        expl_path_info = eval_util.get_generic_path_information(expl_paths)
        logger.record_dict(expl_path_info, prefix="expl/")
        self.writer.add_scalar('expl/Returns Mean', expl_path_info.get('Returns Mean', 0), epoch)
        self.writer.add_scalar('expl/Num Paths', expl_path_info.get('Num Paths', 0), epoch)

        """
        Evaluation
        """
        eval_freq = getattr(self, 'eval_freq', 1)
        if eval_freq > 0 and epoch % eval_freq == 0:
            # path length step 等
            eval_path_info = self.eval_data_collector.get_diagnostics()
            # logger.record_dict(eval_path_info, prefix='eval/')
            self.writer.add_scalar('eval/path length Mean', eval_path_info.get('path length Mean', 0), epoch)
            # xxx
            eval_paths = self.eval_data_collector.get_epoch_paths()
            # if hasattr(self.eval_env, 'get_diagnostics'):
            #     logger.record_dict(self.eval_env.get_diagnostics(eval_paths), prefix='eval/')
            # Returns Reward 等
            eval_path_info = eval_util.get_generic_path_information(eval_paths)
            logger.record_dict(eval_path_info, prefix="eval/")
            self.writer.add_scalar('eval/Returns Mean', eval_path_info.get('Returns Mean', 0), epoch)
            self.writer.add_scalar('eval/Num Paths', eval_path_info.get('Num Paths', 0), epoch)

        """
        Misc
        """
        gt.stamp('logging')
        logger.record_dict(_get_epoch_timings())
        logger.record_tabular('Epoch', epoch)
        logger.dump_tabular(with_prefix=False, with_timestamp=False)

    @abc.abstractmethod
    def training_mode(self, mode):
        """
        Set training mode to `mode`.
        :param mode: If True, training will happen (e.g. set the dropout
        probabilities to not all ones).
        """
        pass
